{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TensorFlow version: 2.0.0\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "import datetime, os\n",
    "#hide tf logs\n",
    "os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'  # or any {'0', '1', '2'},\n",
    "#0 (default) shows all, 1 to filter out INFO logs, 2 to additionally filter out WARNING logs, and 3 to additionally filter out ERROR logs\n",
    "import scipy.optimize\n",
    "import scipy.io\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.ticker\n",
    "import time\n",
    "from pyDOE import lhs         #Latin Hypercube Sampling\n",
    "import pandas as pd\n",
    "import seaborn as sns \n",
    "import codecs, json\n",
    "\n",
    "# generates same random numbers each time\n",
    "np.random.seed(1234)\n",
    "tf.random.set_seed(1234)\n",
    "\n",
    "print(\"TensorFlow version: {}\".format(tf.__version__))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# *Data Prep*\n",
    "\n",
    "Training and Testing data is prepared from the solution file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_1 = np.linspace(-1,1,256)  # 256 points between -1 and 1 [256x1]\n",
    "x_2 = np.linspace(1,-1,256)  # 256 points between 1 and -1 [256x1]\n",
    "\n",
    "X, Y = np.meshgrid(x_1,x_2) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test Data\n",
    "\n",
    "We prepare the test data to compare against the solution produced by the PINN."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_u_test = np.hstack((X.flatten(order='F')[:,None], Y.flatten(order='F')[:,None]))\n",
    "\n",
    "# Domain bounds\n",
    "lb = np.array([-1, -1]) #lower bound\n",
    "ub = np.array([1, 1])  #upper bound\n",
    "\n",
    "a_1 = 1 \n",
    "a_2 = 4\n",
    "\n",
    "usol = np.sin(a_1 * np.pi * X) * np.sin(a_2 * np.pi * Y) #solution chosen for convinience  \n",
    "\n",
    "u = usol.flatten('F')[:,None] "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def trainingdata(N_u,N_f):\n",
    "    \n",
    "    leftedge_x = np.hstack((X[:,0][:,None], Y[:,0][:,None]))\n",
    "    leftedge_u = usol[:,0][:,None]\n",
    "    \n",
    "    rightedge_x = np.hstack((X[:,-1][:,None], Y[:,-1][:,None]))\n",
    "    rightedge_u = usol[:,-1][:,None]\n",
    "    \n",
    "    topedge_x = np.hstack((X[0,:][:,None], Y[0,:][:,None]))\n",
    "    topedge_u = usol[0,:][:,None]\n",
    "    \n",
    "    bottomedge_x = np.hstack((X[-1,:][:,None], Y[-1,:][:,None]))\n",
    "    bottomedge_u = usol[-1,:][:,None]\n",
    "    \n",
    "    all_X_u_train = np.vstack([leftedge_x, rightedge_x, bottomedge_x, topedge_x])\n",
    "    all_u_train = np.vstack([leftedge_u, rightedge_u, bottomedge_u, topedge_u])  \n",
    "     \n",
    "    #choose random N_u points for training\n",
    "    idx = np.random.choice(all_X_u_train.shape[0], N_u, replace=False) \n",
    "    \n",
    "    X_u_train = all_X_u_train[idx, :] #choose indices from  set 'idx' (x,t)\n",
    "    u_train = all_u_train[idx,:]      #choose corresponding u\n",
    "    \n",
    "    '''Collocation Points'''\n",
    "\n",
    "    # Latin Hypercube sampling for collocation points \n",
    "    # N_f sets of tuples(x,t)\n",
    "    X_f = lb + (ub-lb)*lhs(2,N_f) \n",
    "    X_f_train = np.vstack((X_f, X_u_train)) # append training points to collocation points \n",
    "    \n",
    "    return X_f_train, X_u_train, u_train \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PINN \n",
    "\n",
    "$W \\in \\mathcal{R}^{n_{l-1}\\times{n_l}}$ \n",
    "\n",
    "Creating sequential layers using the $\\textit{class}$ tf.Module"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Sequentialmodel(tf.Module): \n",
    "    def __init__(self, layers, name=None):\n",
    "       \n",
    "        self.W = []  #Weights and biases\n",
    "        self.parameters = 0 #total number of parameters\n",
    "        \n",
    "        for i in range(len(layers)-1):\n",
    "            \n",
    "            input_dim = layers[i]\n",
    "            output_dim = layers[i+1]\n",
    "            \n",
    "            #Xavier standard deviation \n",
    "            std_dv = np.sqrt((2.0/(input_dim + output_dim)))\n",
    "\n",
    "            #weights = normal distribution * Xavier standard deviation + 0\n",
    "            w = tf.random.normal([input_dim, output_dim], dtype = 'float64') * std_dv\n",
    "                        \n",
    "            w = tf.Variable(w, trainable=True, name = 'w' + str(i+1))\n",
    "\n",
    "            b = tf.Variable(tf.cast(tf.zeros([output_dim]), dtype = 'float64'), trainable = True, name = 'b' + str(i+1))\n",
    "                    \n",
    "            self.W.append(w)\n",
    "            self.W.append(b)\n",
    "            \n",
    "            self.parameters +=  input_dim * output_dim + output_dim\n",
    "    \n",
    "    def evaluate(self,x):\n",
    "        \n",
    "        #preprocessing input \n",
    "        x = (x - lb)/(ub - lb) #feature scaling\n",
    "        \n",
    "        a = x\n",
    "        \n",
    "        for i in range(len(layers)-2):\n",
    "            \n",
    "            z = tf.add(tf.matmul(a, self.W[2*i]), self.W[2*i+1])\n",
    "            a = tf.nn.tanh(z)\n",
    "            \n",
    "        a = tf.add(tf.matmul(a, self.W[-2]), self.W[-1]) # For regression, no activation to last layer\n",
    "        \n",
    "        return a\n",
    "    \n",
    "    def get_weights(self):\n",
    "\n",
    "        parameters_1d = []  # [.... W_i,b_i.....  ] 1d array\n",
    "        \n",
    "        for i in range (len(layers)-1):\n",
    "            \n",
    "            w_1d = tf.reshape(self.W[2*i],[-1])   #flatten weights \n",
    "            b_1d = tf.reshape(self.W[2*i+1],[-1]) #flatten biases\n",
    "            \n",
    "            parameters_1d = tf.concat([parameters_1d, w_1d], 0) #concat weights \n",
    "            parameters_1d = tf.concat([parameters_1d, b_1d], 0) #concat biases\n",
    "        \n",
    "        return parameters_1d\n",
    "        \n",
    "    def set_weights(self,parameters):\n",
    "                \n",
    "        for i in range (len(layers)-1):\n",
    "\n",
    "            shape_w = tf.shape(self.W[2*i]).numpy() # shape of the weight tensor\n",
    "            size_w = tf.size(self.W[2*i]).numpy() #size of the weight tensor \n",
    "            \n",
    "            shape_b = tf.shape(self.W[2*i+1]).numpy() # shape of the bias tensor\n",
    "            size_b = tf.size(self.W[2*i+1]).numpy() #size of the bias tensor \n",
    "                        \n",
    "            pick_w = parameters[0:size_w] #pick the weights \n",
    "            self.W[2*i].assign(tf.reshape(pick_w,shape_w)) # assign  \n",
    "            parameters = np.delete(parameters,np.arange(size_w),0) #delete \n",
    "            \n",
    "            pick_b = parameters[0:size_b] #pick the biases \n",
    "            self.W[2*i+1].assign(tf.reshape(pick_b,shape_b)) # assign \n",
    "            parameters = np.delete(parameters,np.arange(size_b),0) #delete \n",
    "\n",
    "            \n",
    "    def loss_BC(self,x,y):\n",
    "\n",
    "        loss_u = tf.reduce_mean(tf.square(y-self.evaluate(x)))\n",
    "        return loss_u\n",
    "\n",
    "    def loss_PDE(self, x_to_train_f):\n",
    "    \n",
    "        g = tf.Variable(x_to_train_f, dtype = 'float64', trainable = False)\n",
    "\n",
    "        k = 1    \n",
    "\n",
    "        x_1_f = g[:,0:1]\n",
    "        x_2_f = g[:,1:2]\n",
    "\n",
    "        with tf.GradientTape(persistent=True) as tape:\n",
    "\n",
    "            tape.watch(x_1_f)\n",
    "            tape.watch(x_2_f)\n",
    "\n",
    "            g = tf.stack([x_1_f[:,0], x_2_f[:,0]], axis=1)\n",
    "\n",
    "            u = self.evaluate(g)\n",
    "            u_x_1 = tape.gradient(u,x_1_f)\n",
    "            u_x_2 = tape.gradient(u,x_2_f)\n",
    "\n",
    "        u_xx_1 = tape.gradient(u_x_1,x_1_f)\n",
    "        u_xx_2 = tape.gradient(u_x_2,x_2_f)\n",
    "\n",
    "        del tape\n",
    "\n",
    "        q = -( (a_1*np.pi)**2 + (a_2*np.pi)**2 - k**2 ) * np.sin(a_1*np.pi*x_1_f) * np.sin(a_2*np.pi*x_2_f)\n",
    "\n",
    "        f = u_xx_1 + u_xx_2 + k**2 * u - q #residual\n",
    "        \n",
    "        loss_f = tf.reduce_mean(tf.square(f))\n",
    "\n",
    "        return loss_f, f\n",
    "    \n",
    "    def loss(self,x,y,g):\n",
    "\n",
    "        loss_u = self.loss_BC(x,y)\n",
    "        loss_f, f_plot = self.loss_PDE(g)\n",
    "\n",
    "        loss = loss_u + loss_f\n",
    "\n",
    "        return loss, loss_u, loss_f \n",
    "    \n",
    "    \n",
    "    def gradients(self,x,y,x_train_f):\n",
    "        \n",
    "        with tf.GradientTape() as tape:\n",
    "            tape.watch(self.W)\n",
    "            loss_val, loss_u, loss_f = self.loss(x, y, x_train_f)\n",
    "            \n",
    "        grads = tape.gradient(loss_val,self.W)\n",
    "            \n",
    "        del tape\n",
    "        \n",
    "        grads_1d = [ ] #store 1d grads\n",
    "        \n",
    "        for i in range (len(layers)-1):\n",
    "\n",
    "            grads_w_1d = tf.reshape(grads[2*i],[-1]) #flatten weights \n",
    "            grads_b_1d = tf.reshape(grads[2*i+1],[-1]) #flatten biases\n",
    "\n",
    "            grads_1d = tf.concat([grads_1d, grads_w_1d], 0) #concat grad_weights \n",
    "            grads_1d = tf.concat([grads_1d, grads_b_1d], 0) #concat grad_biases\n",
    "\n",
    "        return loss_val, loss_u, loss_f, grads, grads_1d\n",
    "    \n",
    "    \n",
    "    def optimizer_callback(self,parameters):\n",
    "               \n",
    "        loss_value, loss_u, loss_f = self.loss(X_u_train, u_train, X_f_train)\n",
    "        \n",
    "        u_pred = self.evaluate(X_u_test)\n",
    "        error_vec = np.linalg.norm((u-u_pred),2)/np.linalg.norm(u,2)\n",
    "        \n",
    "        tf.print(loss_value, loss_u, loss_f, error_vec)       "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# *Loss Function*\n",
    "\n",
    "The loss function consists of two parts:\n",
    "1. **loss_BC**: MSE error of boundary losses\n",
    "2. **loss_PDE**: MSE error of collocation points satisfying the PDE\n",
    "\n",
    "**loss** = loss_BC + loss_PDE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:Calling GradientTape.gradient on a persistent tape inside its context is significantly less efficient than calling it outside the context (it causes the gradient ops to be recorded on the tape, leading to increased CPU and memory usage). Only call GradientTape.gradient inside the context if you actually want to trace the gradient in order to compute higher order derivatives.\n",
      "WARNING:tensorflow:Calling GradientTape.gradient on a persistent tape inside its context is significantly less efficient than calling it outside the context (it causes the gradient ops to be recorded on the tape, leading to increased CPU and memory usage). Only call GradientTape.gradient inside the context if you actually want to trace the gradient in order to compute higher order derivatives.\n",
      "3.6991227497924126\n",
      "2.3939688936874912\n",
      "1.2130400880055074\n",
      "0.49199790481270345\n",
      "Training time: 17.39\n",
      "Test Error: 1.01342\n"
     ]
    }
   ],
   "source": [
    "N_u = 400 #Total number of data points for 'u'\n",
    "N_f = 100000 #Total number of collocation points \n",
    "\n",
    "# Training data\n",
    "X_f_train, X_u_train, u_train = trainingdata(N_u,N_f)\n",
    "\n",
    "layers = np.array([2, 50, 50, 50, 1]) #3 hidden layers\n",
    "\n",
    "PINN = Sequentialmodel(layers)\n",
    "\n",
    "init_params = PINN.get_weights().numpy()\n",
    "\n",
    "start_time = time.time() \n",
    "\n",
    "optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3, beta_1=0.9, beta_2=0.999, epsilon=1e-07)\n",
    "\n",
    "num_epochs = 4e0\n",
    "\n",
    "# batch_size = 128\n",
    "\n",
    "# for epoch in range(int(num_epochs)):\n",
    "    \n",
    "#     idx = np.random.choice(X_u_train.shape[0], batch_size, replace=False)\n",
    "    \n",
    "#     loss_value, loss_u, loss_f, grads = PINN.gradients(X_u_train[idx], u_train[idx], X_f_train[idx])\n",
    "\n",
    "#     optimizer.apply_gradients(zip(grads, PINN.W)) #gradient descent\n",
    "        \n",
    "#     if epoch % 100 == 0:\n",
    "\n",
    "#         u_pred = PINN.evaluate(X_u_test)\n",
    "#         error_vec = np.linalg.norm((u-u_pred),2)/np.linalg.norm(u,2)        # Relative L2 Norm of the error (Vector)\n",
    "\n",
    "#         tf.print(loss_value, loss_u, loss_f, error_vec)\n",
    "\n",
    "for epoch in range(int(num_epochs)):\n",
    "\n",
    "    loss_value, loss_u, loss_f, grads, grads_1d = PINN.gradients(X_u_train, u_train, X_f_train)\n",
    "    \n",
    "    tf.print(tf.linalg.norm(grads_1d,2))\n",
    "\n",
    "    optimizer.apply_gradients(zip(grads, PINN.W)) #gradient descent\n",
    "\n",
    "        \n",
    "elapsed = time.time() - start_time                \n",
    "print('Training time: %.2f' % (elapsed))\n",
    "\n",
    "''' Model Accuracy ''' \n",
    "u_pred = PINN.evaluate(X_u_test)\n",
    "\n",
    "error_vec = np.linalg.norm((u-u_pred),2)/np.linalg.norm(u,2)        # Relative L2 Norm of the error (Vector)\n",
    "print('Test Error: %.5f'  % (error_vec))\n",
    "\n",
    "u_pred = np.reshape(u_pred,(256,256),order='F') \n",
    "\n",
    "# #Plot and save image\n",
    "# plt.pcolor(x_1,x_2,u_pred, cmap = 'jet')\n",
    "# plt.axis('scaled')\n",
    "# plt.colorbar()\n",
    "# plt.savefig('Stiff_Helmholtz.png', dpi = 500)\n",
    "\n",
    "# #Error plot\n",
    "# plt.pcolor(x_1,x_2,usol-u_pred, cmap = 'jet')\n",
    "# plt.axis('scaled')\n",
    "# plt.colorbar()\n",
    "# plt.savefig('Stiff_Helmholtz_error.png', dpi = 500)\n",
    "\n",
    "# #Residual plot\n",
    "# loss_f, f_plot = PINN.loss_PDE(X_u_test)\n",
    "# plt.scatter(X_u_test[:,0:1],X_u_test[:,1:2],c=f_plot, cmap = 'jet')\n",
    "# plt.axis('scaled')\n",
    "# plt.colorbar()\n",
    "# plt.savefig('Stiff_Helmholtz_residual.png', dpi = 500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with tf.GradientTape(persistent=True) as tape:\n",
    "    loss_value, loss_u, loss_f = PINN.loss(X_u_train, u_train, X_f_train)\n",
    "    grad_u = tape.gradient(loss_u, PINN.W) \n",
    "    grad_f = tape.gradient(loss_f, PINN.W)\n",
    "    del tape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Writing gradients to 2 separate JSON files \n",
    "\n",
    "L = len(layers)-1 #number of weights matrices\n",
    "\n",
    "for i in range (L*2):\n",
    "\n",
    "    temp = grad_f[i].numpy().tolist() # nested lists with same data, indices\n",
    "    json.dump(temp, codecs.open(\"Non_Stiff_problem/Gradients/gradients_f\" + str(i) + \".json\", 'w', encoding='utf-8'), separators=(',', ':'), sort_keys=True, indent=0)\n",
    "    \n",
    "    temp = grad_u[i].numpy().tolist() # nested lists with same data, indices\n",
    "    json.dump(temp, codecs.open(\"Non_Stiff_problem/Gradients/gradients_u\" + str(i) + \".json\", 'w', encoding='utf-8'), separators=(',', ':'), sort_keys=True, indent=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {},
   "outputs": [],
   "source": [
    "f_1 = {} #residual non-stiff problem\n",
    "u_1 = {} #boundary non-stiff problem\n",
    "\n",
    "f_2 = {} #residual stiff problem\n",
    "u_2 = {} #boundary stiff problem\n",
    "\n",
    "L = 4\n",
    "\n",
    "#Reading gradients from JSON files and storing in a dict \n",
    "#storing gradients of both weights and biases\n",
    "\n",
    "for i in range (L):\n",
    "    \n",
    "    #weights and biases\n",
    "#     obj_text_w = codecs.open(\"Non_Stiff_problem/Gradients/gradients_f\" + str(2*i) + \".json\", 'r', encoding='utf-8').read()\n",
    "#     obj_text_b = codecs.open(\"Non_Stiff_problem/Gradients/gradients_f\" + str(2*i+1) + \".json\", 'r', encoding='utf-8').read()\n",
    "#     f_1['f'+ str(i)] = np.concatenate((np.array(json.loads(obj_text_w)).flatten('F'), np.array(json.loads(obj_text_b)).flatten('F')))\n",
    "    \n",
    "#     obj_text_w = codecs.open(\"Non_Stiff_problem/Gradients/gradients_u\" + str(2*i) + \".json\", 'r', encoding='utf-8').read()\n",
    "#     obj_text_b = codecs.open(\"Non_Stiff_problem/Gradients/gradients_u\" + str(2*i+1) + \".json\", 'r', encoding='utf-8').read()\n",
    "#     u_1['u'+ str(i)] = np.concatenate((np.array(json.loads(obj_text_w)).flatten('F'), np.array(json.loads(obj_text_b)).flatten('F')))\n",
    "    \n",
    "    obj_text_w = codecs.open(\"Stiff_problem/Gradients/gradients_f\" + str(2*i) + \".json\", 'r', encoding='utf-8').read()\n",
    "    obj_text_b = codecs.open(\"Stiff_problem/Gradients/gradients_f\" + str(2*i+1) + \".json\", 'r', encoding='utf-8').read()\n",
    "    f_2['f'+ str(i)] = np.concatenate((np.array(json.loads(obj_text_w)).flatten('F'), np.array(json.loads(obj_text_b)).flatten('F')))\n",
    "    \n",
    "    obj_text_w = codecs.open(\"Stiff_problem/Gradients/gradients_u\" + str(2*i) + \".json\", 'r', encoding='utf-8').read()\n",
    "    obj_text_b = codecs.open(\"Stiff_problem/Gradients/gradients_u\" + str(2*i+1) + \".json\", 'r', encoding='utf-8').read()\n",
    "    u_2['u'+ str(i)] = np.concatenate((np.array(json.loads(obj_text_w)).flatten('F'), np.array(json.loads(obj_text_b)).flatten('F')))\n",
    "    \n",
    "    \n",
    "#     #only weights\n",
    "#     obj_text_w = codecs.open(\"Non_Stiff_problem/Gradients/gradients_f\" + str(2*i) + \".json\", 'r', encoding='utf-8').read()\n",
    "#     f_1['f'+ str(i)] = np.array(json.loads(obj_text_w)).flatten('F')\n",
    "    \n",
    "#     obj_text_w = codecs.open(\"Non_Stiff_problem/Gradients/gradients_u\" + str(2*i) + \".json\", 'r', encoding='utf-8').read()\n",
    "#     u_1['u'+ str(i)] = np.array(json.loads(obj_text_w)).flatten('F')\n",
    "    \n",
    "#     obj_text_w = codecs.open(\"Stiff_problem/Gradients/gradients_f\" + str(2*i) + \".json\", 'r', encoding='utf-8').read()\n",
    "#     f_2['f'+ str(i)] = np.array(json.loads(obj_text_w)).flatten('F')\n",
    "    \n",
    "#     obj_text_w = codecs.open(\"Stiff_problem/Gradients/gradients_u\" + str(2*i) + \".json\", 'r', encoding='utf-8').read()\n",
    "#     u_2['u'+ str(i)] = np.array(json.loads(obj_text_w)).flatten('F')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'L' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-6-a32bd05a7827>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m      5\u001b[0m \u001b[0mfig\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mplt\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mfigure\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;36m4\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mfigsize\u001b[0m\u001b[1;33m=\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;36m13\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;36m4\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      6\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 7\u001b[1;33m \u001b[1;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mrange\u001b[0m \u001b[1;33m(\u001b[0m\u001b[0mL\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m      8\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      9\u001b[0m     \u001b[0max\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mplt\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0msubplot\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;36m1\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;36m4\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mcnt\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mNameError\u001b[0m: name 'L' is not defined"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<Figure size 936x288 with 0 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#Plot the gradients\n",
    "#seaborn = 0.9.0\n",
    "\n",
    "cnt = 1\n",
    "fig = plt.figure(4, figsize=(13, 4))\n",
    "\n",
    "for i in range (L):\n",
    "    \n",
    "    ax = plt.subplot(1, 4, cnt)\n",
    "    ax.set_title('Layer {}'.format(i + 1))\n",
    "    ax.set_yscale('symlog')\n",
    "    \n",
    "    gradients_res = f_1['f'+ str(i)]\n",
    "    gradients_bcs = u_1['u'+ str(i)]\n",
    "      \n",
    "    sns.distplot(gradients_bcs, hist=False,\n",
    "                 kde_kws={\"shade\": False},\n",
    "                 norm_hist=False, label=r'$\\nabla_{\\theta} \\hat J_{BC}$')\n",
    "    \n",
    "    sns.distplot(gradients_res, hist=False,\n",
    "                 kde_kws={\"shade\": False},\n",
    "                 norm_hist=False, label=r'$\\nabla_{\\theta} \\hat J_{PDE}$')\n",
    "\n",
    "    ax.get_legend().remove()\n",
    "    ax.set_xlim([-3, 3])\n",
    "    ax.set_ylim([0,1e6])\n",
    "    cnt += 1\n",
    "    \n",
    "handles, labels = ax.get_legend_handles_labels()\n",
    "fig.legend(handles, labels, loc=\"upper left\", bbox_to_anchor=(0.35, -0.01),\n",
    "             borderaxespad=0, bbox_transform=fig.transFigure, ncol=2)\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "fig.savefig('Gradient_histogram_stiff_Adam.png', dpi = 500, bbox_inches='tight')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
